# -*- coding: utf-8 -*-

"""Wrapper for some crytographic algorithms in pure Python found on the net.

Every algorithm has it's own license in the source file. The wrapper is
licensed under the MIT license.

"""


import os
from hashlib import sha256, sha384, sha512

import aes256, blowfish, twofish

try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO


try:
    unicode_class = unicode
except NameError:
    unicode_class = str


ALGORITHMS = {}
DEFAULT_SALT = '\x99H\x02Y\x86\xca\x03.sGH\x0fL1&\xbd\x8c\x10\x87\x06'


def register(name, klass):
    if not hasattr(klass, 'encrypt') or not hasattr(klass, 'decrypt'):
        raise TypeError("{0} {1!r} doesn't match crypto interface.".format(
            name, klass))
    ALGORITHMS[name.lower()] = klass


def safe_encode(string):
    if isinstance(string, unicode_class):
        string = string.encode('utf-8', errors='replace')
    return string


class _Cipher(object):

    def __init__(self, password, salt, rounds):
        pw = '{0}${1}'.format(safe_encode(salt), safe_encode(password))
        key = sha512(pw).digest()
        for i in range(rounds):
            key = sha512(pw + key).digest()
        self.password = key


class AES(_Cipher):

    @property
    def passwd(self):
        return bytearray(sha256(self.password).digest())

    def encrypt(self, plaindata):
        fp = StringIO(safe_encode(plaindata))
        out = StringIO()
        aes256._encrypt(fp, self.passwd, out)
        try:
            return out.getvalue()
        finally:
            out.close()

    def decrypt(self, cipherdata):
        fp = StringIO(cipherdata)
        out = StringIO()
        aes256._decrypt(fp, self.passwd, out)
        try:
            return out.getvalue()
        finally:
            out.close()


class Blowfish(_Cipher):

    def __init__(self, password, salt, rounds):
        _Cipher.__init__(self, password, salt, rounds)
        self.cipher = blowfish.Blowfish(self.passwd)

    @property
    def passwd(self):
        return sha384(self.password).digest()

    def encrypt(self, plaindata):
        self.cipher.initCTR()
        encrypted = self.cipher.encryptCTR(safe_encode(plaindata))
        return encrypted

    def decrypt(self, cipherdata):
        self.cipher.initCTR()
        decrypted = self.cipher.decryptCTR(cipherdata)
        return decrypted


class Twofish(_Cipher):

    def __init__(self, password, salt, rounds):
        _Cipher.__init__(self, password, salt, rounds)
        self.cipher = twofish.Twofish(self.password)

    def encrypt(self, plaindata):
        data = safe_encode(plaindata)
        pad_bytes = 16 - len(data) % 16
        if pad_bytes == 1:
            data += chr(1)
        elif pad_bytes > 1:
            b = os.urandom(pad_bytes - 1) + chr(pad_bytes)
            data += b
        else:
            # pad_bytes == 0
            b = os.urandom(15) + chr(16)
            data += b
        encrypted = self.cipher.encrypt(data)
        return encrypted

    def decrypt(self, cipherdata):
        decrypted = self.cipher.decrypt(cipherdata)
        pad_length = ord(decrypted[-1])
        return decrypted[:-pad_length]


register('AES', AES)
register('Blowfish', Blowfish)
register('Twofish', Twofish)


def _encrypt_decrypt(text, password, salt, rounds, algorithm, _direction):
    Cipher = ALGORITHMS[algorithm.lower()]
    cipher = Cipher(password, salt, rounds)
    if _direction == 'enc':
        return cipher.encrypt(text)
    else:
        return cipher.decrypt(text)


def encrypt(data, password, salt=DEFAULT_SALT, rounds=16, algorithm='twofish'):
    """Encrypt `data` with `password` and optional `salt` using
    `algorithm`. Builtin algorithms are AES, Blowfish and Twofish.

    :parameters:
        data : unicode or bytestring
            Data to encode (unicode will be encoded to utf-8).
        password : unicode or bytestring
            Password to use for encryption.
        salt : unicode or bytestring
            Optional salt to use with the password. If not given DEFAULT_SALT
            is used.
        rounds : Integer
            Rounds to secure the password. Most people can leave the default
            of 16 here.
        algorithm : string
            The algorithm to use (builtin or every registered),
            default: twofish.

    :returns: Encrypted bytestring.
    :rtype: bytestring
    """
    return _encrypt_decrypt(data, password, salt, rounds, algorithm, 'enc')


def decrypt(data, password, salt=DEFAULT_SALT, rounds=16, algorithm='twofish'):
    """Decrypts `data` which was encrypted using `password` and
    optional `salt` with `algorithm`. For parameters see `encrypt`.

    :returns: The plain decrypted data.
    :rtype: bytestring
    """
    return _encrypt_decrypt(data, password, salt, rounds, algorithm, 'dec')

